from math import floor

import torch

from algorithms.algorithm import Algorithm
from algorithms.base.nsga2 import NSGA2
from config import device
from individual.population import Population
from operators.mutations.polynomial_mutation import PolynomialMutation
from problems.problem import Problem
from utils.vectors import get_ref_vectors, vector_radian


class LMOCSO(Algorithm):
    """
    %------------------------------- Reference --------------------------------
    % Y. Tian, X. Zheng, X. Zhang, and Y. Jin, Efficient large-scale multi-
    % objective optimization based on a competitive swarm optimizer, IEEE
    % Transactions on Cybernetics, 2020, 50(8): 3696-3708.
    %----------------------------------------- --------------------------------
    """

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs: dict,
                 ):
        """
        初始化算法
        :param problem:问题
        :param pop_size: 种群大小，实际大小会根据产生的权重向量个数决定
        :param kwargs: 其他参数
        """
        self.ref_vectors = get_ref_vectors(pop_size, problem.obj_dim)
        pop_size = self.ref_vectors.shape[0]
        super().__init__(pop_size, problem, kwargs)
        # 初始化粒子的速度（vel）为 0
        self.population = self.problem.init_population(pop_size,
                                                       vel=torch.zeros(pop_size, problem.var_dim, dtype=torch.double))
        self.population.screen_populations(
            self.environmental_selection(self.population, self.problem, self.ref_vectors))

    def each_iteration(self):
        # 得到winner,loser的索引
        winner_index, loser_index = self.get_winner_loser_index()
        # 更新个体
        new_population = self.update_individual(self.population, self.problem, winner_index, loser_index)
        # 将新产生的解加入到种群
        self.population += new_population
        # 环境选择
        next_index = self.environmental_selection(self.population, self.problem, self.ref_vectors)
        self.population = self.population[next_index]

    def get_winner_loser_index(self):
        obj = self.population.get_attr('obj')
        fitness = self.cal_fitness(obj)
        [pop_size, _] = obj.shape
        # 随机生成winner和loser
        if pop_size >= 2:
            rank = torch.randperm(pop_size, device=device)[:floor(pop_size / 2) * 2]
            middle_index = int(rank.size(0) / 2)
            loser_index = rank[:middle_index]
            winner_index = rank[middle_index:]
        else:
            rank = torch.tensor([0, 0], device=device, dtype=torch.int)
            loser_index = rank[0]
            winner_index = rank[1]
        # 将成对中适应度低的winner和适应度高的loser对调
        exchange_index = fitness[loser_index] >= fitness[winner_index]
        winner_index[exchange_index], loser_index[exchange_index] = loser_index[exchange_index], winner_index[
            exchange_index]
        return winner_index, loser_index

    @staticmethod
    def cal_fitness(obj):
        # Calculate the fitness by shift-based density
        [size, _] = obj.shape
        f_max = obj.max(0)[0]
        f_min = obj.min(0)[0]
        obj = (obj - f_min) / (f_max - f_min)
        fitness = torch.zeros(size, dtype=torch.float, device=device)
        for i in range(size):
            temp_obj = obj - obj[i]
            temp_obj[i] = float('inf')
            temp_obj[temp_obj < 0] = 0
            fitness[i] = temp_obj.norm(p=2, dim=1).min()
        return fitness

    @staticmethod
    def update_individual(population, problem, winner_index, loser_index):
        dec = population.dec
        vel = population.get_attr('vel')
        winner = dec[winner_index]
        winner_vel = vel[winner_index]
        loser = dec[loser_index]
        loser_vel = vel[loser_index]
        [pop_size, var_dim] = dec.shape
        if pop_size == 1:
            win_size = 1
            winner = winner.unsqueeze(dim=0)
            winner_vel = winner_vel.unsqueeze(dim=0)
        else:
            [win_size, _] = winner.shape
        # 计算随机数
        r0 = torch.rand(win_size, 1, device=device).repeat(1, var_dim)
        r1 = torch.rand(win_size, 1, device=device).repeat(1, var_dim)
        # 更新失败个体
        offspring_vel = r0 * loser_vel + r1 * (winner - loser)
        offspring_dec = loser + offspring_vel + r0 * (offspring_vel - loser_vel)
        # 将更新后的失败个体和winner的dec进行合并
        offspring_dec = torch.cat((offspring_dec, winner), dim=0)
        # 将更新后的失败个体和winner的vel进行合并
        vel = torch.cat((offspring_vel, winner_vel), dim=0)
        # 修正offspring_dec
        offspring_dec = problem.repair_decision(offspring_dec)
        # 变异offspring_dec
        offspring_dec = PolynomialMutation.do(offspring_dec, problem)
        return problem.estimate_population(offspring_dec, vel=vel)

    @staticmethod
    def environmental_selection(population: Population, problem: Problem, ref_vectors):
        obj = population.obj
        [obj_size, M] = obj.shape
        next_index = torch.zeros(obj_size, device=device, dtype=torch.bool)
        # 非支配排序
        front_no, _ = NSGA2.ND_sort(obj, 1)
        obj = obj[front_no == 1]
        next_temp = torch.nonzero(front_no == 1, as_tuple=True)[0]
        # 格式化y
        obj_normal = obj - torch.min(obj, dim=0)[0]
        # 计算参考向量之间的夹角
        ref_angle_radians = vector_radian(ref_vectors, ref_vectors)
        # 将参考向量自己到自己的夹角设置为inf
        ref_angle_radians[torch.eye(ref_angle_radians.shape[0],
                                    ref_angle_radians.shape[1], dtype=torch.bool, device=device)] = float('inf')
        # 记录每个参考向量到他最近参考向量的弧度值
        ref_angle_radians_value = ref_angle_radians.min(1)[0]
        # 计算y_normal到ref_vectors之间的夹角
        obj_to_ref_min = vector_radian(obj_normal, ref_vectors).min(1)
        # 保存y_normal到理他最近的ref_vectors的索引
        obj_to_ref_min_index = obj_to_ref_min[1]
        # 保存y_normal到理他最近的ref_vectors的夹角
        obj_to_ref_min_value = obj_to_ref_min[0]
        # 对聚合后的参考向量去重
        ref_index = torch.unique(obj_to_ref_min_index)
        alpha = (problem.fun_eval / problem.max_fun_eval) ** 2
        # 从每个参考向量中选择apd值最小的个体
        for i in ref_index:
            temp_index = torch.nonzero(obj_to_ref_min_index == i, as_tuple=True)[0]
            APD = (1 + M * alpha * obj_to_ref_min_value[temp_index] / ref_angle_radians_value[i]) * \
                  obj_normal[temp_index].norm(p=2, dim=1)
            min_apd_index = APD.argmin()
            next_index[next_temp[temp_index[min_apd_index]]] = True
        return next_index
